/**
 * @file MultigridVector.h
 * @brief The class to store all the knots of all the grids.
 * @author XDDDD
 * @version 
 * @date 2021-06-10
 */

#ifndef __PAISLEYPARK__MULTIGRIDVECTOR_H__
#define __PAISLEYPARK__MULTIGRIDVECTOR_H__

#include <vector>
#include <iostream>

template <size_t DIM>
class MultigridVector : public std::vector<double> {
public:
	/**
	 * @brief Default construction
	 */
	MultigridVector();
	/**
	 * @brief Construction
	 *
	 * @param _u The given initial guess vector.
	 * @param _N The given order. The length of _u is 2^N.
	 */
	MultigridVector(const std::vector<double>& _u, const size_t _N);
	/**
	 * @brief Set the Vector.
	 *
	 * @param _u The given initial guess vector.
	 * @param _N The given order.
	 */
	void set(const std::vector<double>& _u, const size_t _N);
	/**
	 * @brief Return the value of one knots.
	 *
	 * @param _coord The given coordinate of the knot. 
	 * @param _M The order of the grid.
	 *
	 * @return The value of the knot.
	 */
	const double& operator()(const std::vector<size_t>& _coord, const size_t _M) const;
	/**
	 * @brief Return the value of one knot.
	 *
	 * @param _coord The given coordinate of the knot. 
	 * @param _M The order of the grid.
	 *
	 * @return The value of the knot.
	 */
	double& operator()(const std::vector<size_t>& _coord, const size_t _M);
	/**
	 * @brief Return the value of one knot.
	 *
	 * @param index The index of the knot.
	 * @param _M The order of the grid.
	 *
	 * @return The value of the knot
	 */
	const double& operator()(const size_t index, const size_t _M) const;
	/**
	 * @brief Return the value of one knot.
	 *
	 * @param index The index of the knot.
	 * @param _M The order of the grid.
	 *
	 * @return The value of the knot
	 */
	double& operator()(const size_t index, const size_t _M);
	/**
	 * @brief Get the order of the finest grid.
	 *
	 * @return  The order of the finest grid.
	 */
	const size_t get_N() const;
	/**
	 * @brief Get the start index of the grid.
	 *
	 * @param _M The order of the grid we want.
	 *
	 * @return The start index of the grid with order _M
	 */
	const size_t get_start_index(const size_t _M) const;
	/**
	 * @brief Set all the elements in one grid to zero.
	 *
	 * @param _M The order of the grid
	 */
	void clean(const size_t _M);
	/**
	 * @brief Set all the elements except the boundary one in one grid to zero.
	 *
	 * @param _M The order of the grid,
	 */
	void cleanpart(const size_t _M);
private:
	/**
	 * @brief The order of the finest grid.
	 */
	size_t __N;
	/**
	 * @brief The vector to store the starting index of each grid.
	 */
	std::vector<double> __start_index;
};

template <size_t DIM>
MultigridVector<DIM>::MultigridVector() : std::vector<double>() {
	__N = 0;
	__start_index.clear();
};

template <size_t DIM>
MultigridVector<DIM>::MultigridVector(const vector<double>& _u, const size_t _N) : std::vector<double>() {
	__N = _N;
	size_t l = 0;
	size_t index = 0;
	__start_index.clear();
	__start_index.resize(_N);
	size_t size = 0;
	clear();
	for(size_t i = 0; i != _N; i++) {
		index = index + l;
		__start_index[i] = index;
		l = 1;
		for(size_t j = 0; j!= DIM; j++) {
			l = l*(1 << i + 1) + l;
		}
		size += l;
	}
	resize(size);
	for(size_t i = 0; i != l; i++) {
		(*this)[__start_index[_N - 1] + i] = _u[i];
	}
};

template <size_t DIM>
void MultigridVector<DIM>::set(const vector<double>& _u, const size_t _N) {
	__N = _N;
	size_t l = 0;
	size_t index = 0;
	__start_index.resize(_N);
	size_t size = 0;
	clear();
	for(size_t i = 0; i != _N; i++) {
		index = index + l;
		__start_index[i] = index;
		l = 1;
		for(size_t j = 0; j!= DIM; j++) {
			l = l*(1 << i + 1) + l;
		}
		size += l;
	}

	resize(size);	
	for(size_t i = 0; i != l; i++) {
		(*this)[__start_index[_N - 1] + i] = _u[i];
	}
};

template <size_t DIM>
const double& MultigridVector<DIM>::operator()(const std::vector<size_t>& _coord, const size_t _M) const {
	if(_M > __N) {
		std::cout << "Out of range!" << std::endl;
		std::exit(-1);
	}
	if(_coord.size() != DIM) {
		std::cout << "The wrong dimension of the given coordinate." << std::endl;
		std::exit(-1);
	}
	size_t index = __start_index[_M - 1];
	size_t dl = 1;
	for(size_t i = 0; i != DIM; i++) {
		index = index + _coord[i]*dl;
		dl = dl*(1 << _M) + dl;
	}
	return (*this)[index];
};

template <size_t DIM>
double& MultigridVector<DIM>::operator()(const std::vector<size_t>& _coord, const size_t _M) {
	if(_M > __N) {
		std::cout << "Out of range!" << std::endl;
		std::exit(-1);
	}
	if(_coord.size() != DIM) {
		std::cout << "The wrong dimension of the given coordinate." << std::endl;
		std::exit(-1);
	}
	size_t index = __start_index[_M - 1];
	size_t dl = 1;
	for(size_t i = 0; i != DIM; i++) {
		index = index + _coord[i]*dl;
		dl = dl*(1 << _M) + dl;
	}
	return (*this)[index];
};

template <size_t DIM>
const double& MultigridVector<DIM>::operator()(const size_t index, const size_t _M) const {
	return (*this)[(*this).get_start_index(_M) + index];
};

template <size_t DIM>
double& MultigridVector<DIM>::operator()(const size_t index, const size_t _M) {
	return (*this)[(*this).get_start_index(_M) + index];
};

template <size_t DIM>	
const size_t MultigridVector<DIM>::get_N() const {
	return __N;
};

template <size_t DIM>
const size_t MultigridVector<DIM>::get_start_index(const size_t _M) const {
	return __start_index[_M - 1];
};

template <size_t DIM>
void MultigridVector<DIM>::clean(const size_t _M) {
	size_t num_1d = (1 << _M) + 1;
	size_t l = 1;
	for(size_t i = 0; i != DIM; i++) {
		l *= num_1d;
	}
	for(size_t i = 0; i != l; i++) {
		(*this)[__start_index[_M - 1] + i] = 0;
	}
};

template <size_t DIM>
void MultigridVector<DIM>::cleanpart(const size_t _M) {
	std::vector<size_t> coord;
	coord.clear();
	coord.resize(DIM);
	size_t pole = 0;
	size_t num_1d = (1 << _M) + 1;
	for(size_t i = 0; i != DIM; i++) {
		coord[i] = 1;
	}

	while(pole != DIM) {
		pole = 0;
		(*this)(coord, _M) = 0;
		coord[pole]++;
		while(coord[pole] == num_1d - 1) {
			coord[pole] = 1;
			pole++;
			if(pole == DIM)
				break;
			coord[pole]++;
		}
	}
};
#else
//Do nothing.
#endif
